#include "mem.h"
#include "cpu/difftest.h"
#include "model.h"
#include "ioe/ioe.h"
#include "device/device.h"
#include "utils/systime.h"
#include "cpu/cpu.h"

MTracer mtracer;

static uint64_t *pmem;

bool is_device_addr(long long addr)
{
    if ((addr & 0xf0000000) == 0xa0000000)
        return true;
    if ((addr & 0xf0000000) == 0x00000000)
        return true;
    return false;
}

void device_write(long long waddr, long long wdata, char wmask)
{
    // difftest_skip_ref();
    if (waddr == SERIAL_PORT)
    {
        putchar((char)wdata);
        fflush(stdout);
        return;
    }
    if (waddr == VGACTL_ADDR)
    {
        return;
    }
    if (waddr == VGACTL_ADDR + 4)
    {
        update_screen();
        return;
    }
    if (waddr >= FB_ADDR && waddr < FB_ADDR + sizeof(frame_buff))
    {
        uint64_t paddr = (waddr - FB_ADDR) / sizeof(uint64_t);
        uint64_t pdata = 0;

        uint64_t mdata = ((uint64_t *)frame_buff)[paddr];
        for (int i = 0; i < sizeof(uint64_t); i++)
        {
            if (BITS(wmask, i, i))
            {
                pdata |= BITS(wdata, i * 8 + 7, i * 8) << i * 8;
            }
            else
            {
                pdata |= BITS(mdata, i * 8 + 7, i * 8) << i * 8;
            }
        }
        ((uint64_t *)frame_buff)[paddr] = pdata;
        return;
    }
    if (waddr >= FB_ADDR && waddr < FB_ADDR + sizeof(frame_buff))
        return;
    LogErr("device: waddr:0x%016llx wdata:0x%016llx wmask:0x%x", waddr, wdata, wmask);
    continue_flag = false;
}

long long device_read(long long raddr, char rsize, char is_icache)
{
    // difftest_skip_ref();
    if (raddr == RTC_ADDR)
    {
        // static uint64_t start = 0;
        // uint64_t us = SYS_GetTicks();
        // if (start == 0)
        //     start = us;
        // return us - start;
        return total_cycles;
    }
    if (raddr == VGACTL_ADDR)
    {
        return (VGA_W << 16) | VGA_H;
    }
    if (raddr == SERIAL_PORT)
    {
        return 0;
    }
    if (raddr == KBD_ADDR)
    {
        return key_dequeue();
    }

    // if (raddr >= FB_ADDR && raddr < FB_ADDR + sizeof(frame_buff))
    // {
    //     uint64_t paddr = (raddr - FB_ADDR);
    //     uint64_t rdata = 0;
    //     uint8_t *p1;
    //     uint16_t *p2;
    //     uint32_t *p4;
    //     uint64_t *p8;
    //     switch (rsize)
    //     {
    //     case 1:
    //         paddr /= 1;
    //         p1 = (uint8_t *)pmem;
    //         rdata = p1[paddr];
    //         break;
    //     case 2:
    //         paddr /= 2;
    //         p2 = (uint16_t *)pmem;
    //         rdata = p2[paddr];
    //         break;
    //     case 4:
    //         paddr /= 4;
    //         p4 = (uint32_t *)pmem;
    //         rdata = p4[paddr];
    //         break;
    //     case 8:
    //         paddr /= 8;
    //         p8 = (uint64_t *)pmem;
    //         rdata = p8[paddr];
    //         break;

    //     default:
    //         assert(0);
    //         break;
    //     }
    //     return rdata;
    // }
    LogErr("device: raddr:0x%016llx rsize:0x%x", raddr, rsize);
    continue_flag = false;
}

void memory_out_of_bond(uint64_t addr)
{
    LogErr("address = " FMT_PADDR " is out of bound of pmem [" FMT_PADDR ", " FMT_PADDR "] at pc = " FMT_WORD "\n",
           addr, PMEM_LEFT, PMEM_RIGHT, -1);
}
long long pmem_read(long long raddr, char rsize, char is_icache)
{
    if (is_device_addr(raddr))
    {
        return device_read(raddr, rsize, 0);
    }

    if (raddr < CONFIG_MBASE || raddr - CONFIG_MBASE >= CONFIG_MSIZE)
    {
        memory_out_of_bond(raddr);
        continue_flag = false;
        return 0;
    }
    if (is_icache)
    {
        uint32_t *icahce = (uint32_t *)pmem;
        uint64_t paddr = (raddr - CONFIG_MBASE) / sizeof(uint32_t);
        uint32_t inst = icahce[paddr];
        return inst;
    }
    else
    {
        uint64_t paddr = (raddr - CONFIG_MBASE);
        uint64_t rdata = 0;
        uint8_t *p1;
        uint16_t *p2;
        uint32_t *p4;
        uint64_t *p8;
        switch (rsize)
        {
        case 1:
            paddr /= 1;
            p1 = (uint8_t *)pmem;
            rdata = p1[paddr];
            break;
        case 2:
            paddr /= 2;
            p2 = (uint16_t *)pmem;
            rdata = p2[paddr];
            break;
        case 4:
            paddr /= 4;
            p4 = (uint32_t *)pmem;
            rdata = p4[paddr];
            break;
        case 8:
            paddr /= 8;
            p8 = (uint64_t *)pmem;
            rdata = p8[paddr];
            break;

        default:
            assert(0);
            break;
        }
        // Log("paddr:0x%016lx ",paddr);

        IFDEF(CONFIG_MTRACE, MTracer_Data data; data.addr = raddr; data.data = rdata; data.size = rsize; data.type = LOAD; mtracer.save(data));
        Log("raddr:0x%016llx rdata:0x%016lx \n", raddr, rdata);
        return rdata;
    }
}
void pmem_write(long long waddr, long long wdata, char wmask)
{
    if (is_device_addr(waddr))
    {
        device_write(waddr, wdata, wmask);
        return;
    }
    Log("waddr:0x%016llx wdata:0x%016llx wmask:%d\n", waddr, wdata, wmask);
    IFONE(CONFIG_MTRACE, MTracer_Data data; data.addr = waddr; data.data = wdata; data.mask = wmask; data.type = STORE; mtracer.save(data));
    if (waddr < CONFIG_MBASE || waddr - CONFIG_MBASE >= CONFIG_MSIZE)
    {
        printf("write ");
        memory_out_of_bond(waddr);
        return;
    }

    uint64_t paddr = (waddr - CONFIG_MBASE) / sizeof(uint64_t);
    uint64_t pdata = 0;
    uint64_t mdata = pmem[paddr];
    for (int i = 0; i < sizeof(uint64_t); i++)
    {
        if (BITS(wmask, i, i))
        {
            pdata |= BITS(wdata, i * 8 + 7, i * 8) << i * 8;
        }
        else
        {
            pdata |= BITS(mdata, i * 8 + 7, i * 8) << i * 8;
        }
    }
    pmem[paddr] = pdata;
}

long long mem_read(long long raddr, svBit is_icache)
{
    return pmem_read(raddr, 8, is_icache);
}
void mem_write(long long waddr, long long wdata, char wmask)
{
    pmem_write(waddr, wdata, wmask);
}

long long axi4_read(int raddr, char rsize)
{
    uintptr_t paddr = (uint32_t)raddr;
    if (is_device_addr(paddr))
    {
        return device_read(paddr, rsize, 0);
    }

    if (paddr < CONFIG_MBASE || paddr - CONFIG_MBASE >= CONFIG_MSIZE)
    {
        memory_out_of_bond(paddr);
        continue_flag = false;
        return 0;
    }
    paddr -= CONFIG_MBASE;
    // printf("axi4_read: raddr:0x%08x rsize:%d\n",raddr,rsize);
    uint8_t *p8 = (uint8_t *)((uintptr_t)paddr + (uintptr_t)pmem);
    uint16_t *p16 = (uint16_t *)((uintptr_t)paddr + (uintptr_t)pmem);
    uint32_t *p32 = (uint32_t *)((uintptr_t)paddr + (uintptr_t)pmem);
    uint64_t *p64 = (uint64_t *)((uintptr_t)paddr + (uintptr_t)pmem);
    union
    {
        uint8_t d8[8];
        uint16_t d16[4];
        uint32_t d32[2];
        uint64_t d64[1];
        uint64_t val;
    } res;
    switch (rsize)
    {
    case 0:
        memset(res.d8, *p8, 8);
        break;
    case 1:
        res.d16[0] = res.d16[1] = res.d16[2] = res.d16[3] = *p16;
        break;
    case 2:
        res.d32[0] = res.d32[1] = *p32;
        break;
    case 3:
        res.d64[0] = *p64;
        break;
    default:
        assert(0);
    }
    return res.val;
}

void axi4_device_write(long long waddr, long long wdata, char wsize)
{
    if (waddr == SERIAL_PORT)
    {
        putchar((char)wdata);
        return;
    }
    if (waddr == VGACTL_ADDR)
    {
        return;
    }
    if (waddr == VGACTL_ADDR + 4)
    {
        update_screen();
        return;
    }
    if (waddr >= FB_ADDR && waddr < FB_ADDR + sizeof(frame_buff))
    {
        uint64_t paddr = (waddr - FB_ADDR);

        uint8_t *p8 = (uint8_t *)((uintptr_t)paddr + (uintptr_t)frame_buff);
        uint16_t *p16 = (uint16_t *)((uintptr_t)paddr + (uintptr_t)frame_buff);
        uint32_t *p32 = (uint32_t *)((uintptr_t)paddr + (uintptr_t)frame_buff);
        uint64_t *p64 = (uint64_t *)((uintptr_t)paddr + (uintptr_t)frame_buff);
        switch (wsize)
        {
        case 0:
            *p8 = wdata;
            break;
        case 1:
            *p16 = wdata;
            break;
        case 2:
            *p32 = wdata;
            break;
        case 3:
            *p64 = wdata;
            break;
        default:
            assert(0);
        }
        return;
    }
    LogErr("device: waddr:0x%016llx wdata:0x%016llx wsize:0x%x", waddr, wdata, wsize);
    continue_flag = false;
}
void axi4_write(int waddr, long long wdata, char wsize)
{
    uintptr_t paddr = (uint32_t)waddr;
    if (is_device_addr(paddr))
    {
        return axi4_device_write(paddr, wdata, wsize);
    }

    if (paddr < CONFIG_MBASE || paddr - CONFIG_MBASE >= CONFIG_MSIZE)
    {
        memory_out_of_bond(paddr);
        continue_flag = false;
        return;
    }
    paddr -= CONFIG_MBASE;
    // printf("axi4_read: raddr:0x%08x rsize:%d\n",raddr,rsize);
    uint8_t *p8 = (uint8_t *)((uintptr_t)paddr + (uintptr_t)pmem);
    uint16_t *p16 = (uint16_t *)((uintptr_t)paddr + (uintptr_t)pmem);
    uint32_t *p32 = (uint32_t *)((uintptr_t)paddr + (uintptr_t)pmem);
    uint64_t *p64 = (uint64_t *)((uintptr_t)paddr + (uintptr_t)pmem);
    switch (wsize)
    {
    case 0:
        *p8 = wdata;
        break;
    case 1:
        *p16 = wdata;
        break;
    case 2:
        *p32 = wdata;
        break;
    case 3:
        *p64 = wdata;
        break;
    default:
        assert(0);
    }
}
uint8_t *guest_to_host(uint64_t paddr)
{
    return (uint8_t *)(pmem + paddr - CONFIG_MBASE);
}
void deinit_mem()
{
    free(pmem);
}
void init_mem()
{
    pmem = (uint64_t *)malloc(CONFIG_MSIZE);
    atexit(deinit_mem);
    uint32_t *p = (uint32_t *)pmem;
    int i;
    for (i = 0; i < (int)(CONFIG_MSIZE / sizeof(p[0])); i++)
    {
        p[i] = rand();
    }
}

static const uint32_t img[] = {
    0x00000297, // auipc t0,0
    0x0002b823, // sd  zero,16(t0)
    0x0102b503, // ld  a0,16(t0)
    0x00100073, // ebreak (used as nemu_trap)
    0xdeadbeef, // some data
};

long load_img(char *img_file)
{
    if (img_file == NULL)
    {
        Log("No image is given. Use the default build-in image.\n");
        memcpy(guest_to_host(RESET_VECTOR), img, sizeof(img));
        return 4096; // built-in image size
    }
    else
    {
        Log("load img : %s", img_file);
    }

    FILE *fp = fopen(img_file, "rb");
    assert(fp);

    fseek(fp, 0, SEEK_END);
    long size = ftell(fp);

    Log("The image is %s, size = %ld", img_file, size);

    fseek(fp, 0, SEEK_SET);
    int ret = fread(guest_to_host(RESET_VECTOR), size, 1, fp);
    assert(ret == 1);

    fclose(fp);
    return size;
}
